import socket
import struct
import random
import logging
from typing import Tuple, Optional

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# STUN消息类型
BINDING_REQUEST = 0x0001
BINDING_RESPONSE = 0x0101

# STUN属性类型
XOR_MAPPED_ADDRESS = 0x0020
MAGIC_COOKIE = 0x2112A442

class StunServer:
    def __init__(self, host: str = '0.0.0.0', port: int = 3478):
        """初始化STUN服务器

        Args:
            host (str): 监听地址
            port (int): 监听端口
        """
        self.host = host
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind((host, port))
        logger.info(f'STUN服务器启动在 {host}:{port}')

    def run(self):
        """运行STUN服务器主循环"""
        while True:
            try:
                data, addr = self.sock.recvfrom(1024)
                self.handle_request(data, addr)
            except Exception as e:
                logger.error(f'处理请求时发生错误: {e}')

    def handle_request(self, data: bytes, addr: Tuple[str, int]):
        """处理STUN请求

        Args:
            data (bytes): 接收到的数据
            addr (Tuple[str, int]): 客户端地址
        """
        try:
            logger.info(f'收到来自 {addr} 的请求，数据长度: {len(data)} 字节')
            msg_type = struct.unpack('!H', data[0:2])[0]
            logger.debug(f'消息类型: 0x{msg_type:04x}')
            if msg_type == BINDING_REQUEST:
                transaction_id = data[8:20]
                logger.info(f'收到Binding请求，事务ID: {transaction_id.hex()}')
                self.send_binding_response(addr, transaction_id)
                logger.info(f'已处理来自 {addr} 的Binding请求')
            else:
                logger.warning(f'收到未知消息类型: 0x{msg_type:04x}')
        except Exception as e:
            logger.error(f'解析请求时发生错误: {e}')

    def send_binding_response(self, addr: Tuple[str, int], transaction_id: bytes):
        """发送Binding响应

        Args:
            addr (Tuple[str, int]): 客户端地址
            transaction_id (bytes): 事务ID
        """
        logger.debug(f'开始构建Binding响应，目标地址: {addr}')
        # 构建响应头
        header = struct.pack('!HHI', 
            BINDING_RESPONSE,  # 消息类型
            8,                # 消息长度（不包括头部）
            MAGIC_COOKIE      # Magic Cookie
        ) + transaction_id
        logger.debug(f'响应头: 类型=0x{BINDING_RESPONSE:04x}, 长度=8, Magic Cookie=0x{MAGIC_COOKIE:08x}')

        # 构建XOR-MAPPED-ADDRESS属性
        ip_parts = [int(part) for part in addr[0].split('.')]
        ip_packed = struct.pack('!BBBB', *ip_parts)
        port = addr[1]
        logger.debug(f'原始IP: {addr[0]}, 端口: {port}')
        
        # XOR操作
        xored_port = port ^ (MAGIC_COOKIE >> 16)
        xored_ip = bytes([ip_packed[i] ^ ((MAGIC_COOKIE >> (24 - 8 * i)) & 0xFF) for i in range(4)])
        logger.debug(f'XOR后的端口: {xored_port}, XOR后的IP: {xored_ip.hex()}')

        # 属性头部
        attr_header = struct.pack('!HH', 
            XOR_MAPPED_ADDRESS,  # 属性类型
            8                    # 属性长度
        )

        # 属性值
        attr_value = struct.pack('!HBB', 
            0x0001,             # IPv4地址族
            xored_port >> 8,    # XOR后的端口（高字节）
            xored_port & 0xFF   # XOR后的端口（低字节）
        ) + xored_ip

        # 发送响应
        response = header + attr_header + attr_value
        logger.debug(f'发送响应，总长度: {len(response)} 字节')
        self.sock.sendto(response, addr)
        logger.info(f'已发送Binding响应到 {addr}')

def main():
    server = StunServer()
    try:
        server.run()
    except KeyboardInterrupt:
        logger.info('服务器关闭')

if __name__ == '__main__':
    main()
